package com.example.routelink.gateway.filter;

import cn.dev33.satoken.secure.SaBase64Util;
import cn.dev33.satoken.secure.SaSecureUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.example.routelink.gateway.entity.LogRecord;
import com.example.routelink.gateway.entity.ResultData;
import com.example.routelink.gateway.entity.ResultEnum;
import com.example.routelink.gateway.utils.IpUtils;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.HashMap;
import java.util.Objects;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.*;

@Component
@Slf4j
public class GatewayResponseFilter implements GlobalFilter, Ordered {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        Instant startTime = Instant.now();
        ServerHttpResponse originalResponse = exchange.getResponse();
        // 暂时移除请求头中的Accept-Encoding 后端服务无法直接判断是否支持gzip
//        exchange.getRequest().mutate().header("accept-encoding", "");
        ServerHttpResponseDecorator response = new ServerHttpResponseDecorator(originalResponse) {
            @NotNull
            @Override
            public Mono<Void> writeWith(@NotNull Publisher<? extends DataBuffer> body) {
                Instant endTime = Instant.now();
                LogRecord logInfo = new LogRecord();
                long elapsedTimeMillis = java.time.Duration.between(startTime, endTime).toMillis();
                logInfo.setElapsedTimeMillis(elapsedTimeMillis);
                Object cachedBody = exchange.getAttribute(ServerWebExchangeUtils.CACHED_REQUEST_BODY_ATTR);
                if (cachedBody != null) {
                    logInfo.setReqBodyData(cachedBody.toString());
                }
                String userAgent = exchange.getRequest().getHeaders().getFirst("user-agent");
                logInfo.setUserAgent(userAgent);
                logInfo.setUserIp(IpUtils.getRealIpAddress(exchange.getRequest()));

                String realUrl = exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR).toString();
                logInfo.setRealUrl(realUrl);
                if (Objects.requireNonNull(getStatusCode()).isError()) {
                    log.info("后端服务请求异常 状态码{} 耗时 {}，地址 {}", getStatusCode().value(), elapsedTimeMillis, realUrl);
                    Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                    ResultData resultData = ResultData.builder().code(ResultEnum.ErrUpstreamServer.getCode()).message(ResultEnum.ErrUpstreamServer.getMsg()).build();
                    DataBufferFactory dataBufferFactory = originalResponse.bufferFactory();
                    byte[] resByteArray = JSON.toJSONString(resultData).getBytes(StandardCharsets.UTF_8);
                    originalResponse.getHeaders().setContentLength(resByteArray.length);
                    originalResponse.getHeaders().setContentType(MediaType.APPLICATION_JSON);
                    return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                        return dataBufferFactory.wrap(resByteArray);
                    }));
                }
                if (Objects.equals(getStatusCode(), HttpStatus.OK) && body instanceof Flux) {
                    String originalResponseContentType = exchange.getAttribute(ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR);
                    if (StrUtil.isNotBlank(originalResponseContentType) && originalResponseContentType.contains(MediaType.APPLICATION_JSON_VALUE)) {
                        Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                        return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                            DataBufferFactory dataBufferFactory = originalResponse.bufferFactory();
                            DataBuffer join = dataBufferFactory.join(dataBuffers);
                            byte[] byteArray = new byte[join.readableByteCount()];
                            join.read(byteArray);
                            DataBufferUtils.release(join);
                            String originalResponseBody = new String(byteArray, StandardCharsets.UTF_8);
                            logInfo.setResBodyData(originalResponseBody);
                            Boolean resIsStackTrace = exchange.getAttribute(RewriteStackTraceResponseBodyFilterGatewayFilterFactory.RES_IS_STACK_TRACE);
                            if (resIsStackTrace != null && resIsStackTrace) {
                                if (isStackTrace(originalResponseBody)) {
                                    ResultData resultData = ResultData.builder().code(ResultEnum.ErrUpstreamServer.getCode()).message(ResultEnum.ErrUpstreamServer.getMsg()).build();
                                    originalResponseBody = JSON.toJSONString(resultData);
                                }
                            }
                            Boolean resAttribute = exchange.getAttribute(EncryptResponseBodyFilterGatewayFilterFactory.RES_IS_ENCRYPT);
                            if (resAttribute != null && resAttribute) {
                                String resAttributeKey = exchange.getAttribute(EncryptResponseBodyFilterGatewayFilterFactory.RES_ENCRYPT_KEY);
                                originalResponseBody = encrypt(originalResponseBody, resAttributeKey);
                            }
                            byte[] resByteArray = originalResponseBody.getBytes(StandardCharsets.UTF_8);
                            originalResponse.getHeaders().setContentLength(resByteArray.length);
                            long responseBodyTime = java.time.Duration.between(endTime, Instant.now()).toMillis();
                            logInfo.setResBodyTime(responseBodyTime);
                            logInfo.setResBodySize(originalResponseBody.length());
                            log.info("[过滤器] 上游服务请求记录 {}", JSON.toJSONString(logInfo));
                            return dataBufferFactory.wrap(resByteArray);
                        }));
                    }
                }

                return super.writeWith(body);
            }

            @NotNull
            @Override
            public Mono<Void> writeAndFlushWith(@NotNull Publisher<? extends Publisher<? extends DataBuffer>> body) {
                return writeWith(Flux.from(body).flatMapSequential(p -> p));
            }
        };

        return chain.filter(exchange.mutate().response(response).build());
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE;
    }

    public static String encrypt(Object str, String secret) {
        return SaBase64Util.encode(SaSecureUtil.aesEncrypt(secret, String.valueOf(str)));
    }

    public boolean isStackTrace(String text) {
        // 检查是否包含 "Exception in thread" 或 "Caused by" 关键字
        if (!text.contains("Exception in thread") && !text.contains("Caused by")) {
            return false;
        }
        // 检查是否包含标准的异常类名模式
        if (!text.matches(".*\\b(java\\.lang|javax|org|com)\\..*Exception.*")) {
            return false;
        }
        // 检查是否包含 "at" 后跟类名、方法名和文件名及行号的模式
        String atPattern = "at \\w+\\.\\w+\\(\\w+\\.java:\\d+\\)";
        if (!text.matches(".*" + atPattern + ".*")) {
            return false;
        }
        // 如果所有条件都满足，则认为是有效的异常堆栈信息
        return true;
    }
}